from boundary_conditions import BoundaryCondition

class AcousticswithAdvectionBC(BoundaryCondition):
    def __init__(self, velocity):
        self.velocity = velocity

    def apply_bc(self, grid):
        q = grid.q
        nb = grid.nb
        
        # Wall boundary conditions
        bcells = grid.boundary_cells["wall"]["cells"]
        images = grid.boundary_cells["wall"]["images"]
        
        for i, cell in enumerate(bcells):
            image = images[ i ]
            
            qb = cell.q     # boundary cell
            qi = image.q # image value

            # set the boundary cell to the image value
            qb[:] = qi[:]
            
            # negate velocity for solid wall
            qb[self.velocity] *= -1
            
        # Periodic boundary conditions
        bcells = grid.boundary_cells["periodic"]["cells"]
        images = grid.boundary_cells["periodic"]["images"]
        
        for i, cell in enumerate(bcells):
            image = images[i]
            
            qb = cell.q     # boundary value
            qi = image.q    # image value
            
            # periodic conditions for the second  component
            qb[2] = qi[2]
        
class AcousticswithAdvectionLeft(BoundaryCondition):
    r"""Special boundary conditions for the problem coupling acoustics with
    advection.

    The two problems are decoupled but are solved together to
    demonstrate schemes for linear equations.

    For the acoustics problem, we use reflecting solid wall boundary
    conditions while for the advection problem, we implement a
    periodic boundary condition.

    """
    def __init__(self, velocity):
        """Constructor.

        Parameters:
        ------------

        velocity : int
            Index for the velocity variable in the vector of
            conserved variables.

        """
        self.velocity = velocity

    def apply_bc(self, grid):
        """Apply the boundary condition."""
        q = grid.q
        nvar = grid.nvar
        
        q1 = q[:, 1]; q2 = q[:,2]
        q0 = q[:, 0]; q3 = q[:,3]

        q3m = q[:, -3]; q4m = q[:, -4]

        q1[:] = q2
        q0[:] = q3

        q1[self.velocity] *= -1
        q0[self.velocity] *= -1

        # implement the periodic bc for the third component
        q1[2] = q3m[2]
        q0[2] = q4m[2]

class AcousticswithAdvectionRight(BoundaryCondition):
    r"""Special boundary conditions for the problem coupling acoustics with
    advection.

    The two problems are decoupled but are solved together to
    demonstrate schemes for linear equations.

    For the acoustics problem, we use reflecting solid wall boundary
    conditions while for the advection problem, we implement a
    periodic boundary condition.

    """
    def __init__(self, velocity):
        """Constructor.

        Parameters:
        ------------

        velocity : int
            Index for the velocity variable in the vector of
            conserved variables.

        """

        self.velocity = velocity

    def apply_bc(self, grid):
        """Apply the boundary condition."""
        q = grid.q
        nvar = grid.nvar

        q0 = q[:, -2]; q3 = q[:,-3]
        q1 = q[:, -1]; q2 = q[:,-4]

        q2p = q[:, 2]; q3p = q[:, 3]
        
        q0[:] = q3
        q1[:] = q2

        q0[self.velocity] *= -1
        q1[self.velocity] *= -1

        # implement the periodic boundary condition for the third
        # component
        q0[2] = q2p[2]
        q1[2] = q3p[2]
